import os
import json
import socket
import threading
import hmac
import hashlib
from typing import Dict, Any, Callable, Optional
from dotenv import load_dotenv
from loguru import logger

# 加载环境变量
load_dotenv()

class MCPServer:
    """MCP协议服务器，用于处理客户端请求"""
    
    def __init__(self, request_handler: Callable[[Dict[str, Any]], Dict[str, Any]]):
        """
        初始化MCP服务器
        
        Args:
            request_handler: 请求处理函数，接收请求数据并返回响应数据
        """
        self.host = os.getenv("MCP_HOST", "localhost")
        self.port = int(os.getenv("MCP_PORT", "8080"))
        self.secret = os.getenv("MCP_SECRET", "")
        self.request_handler = request_handler
        self.server_socket = None
        self.running = False
        self.clients = []
        
        if not self.secret:
            logger.warning("MCP服务器未配置密钥，将使用不安全的通信")
    
    def start(self):
        """启动MCP服务器"""
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        
        try:
            self.server_socket.bind((self.host, self.port))
            self.server_socket.listen(5)
            self.running = True
            logger.info(f"MCP服务器已启动，监听 {self.host}:{self.port}")
            
            # 启动接受连接的线程
            accept_thread = threading.Thread(target=self._accept_connections)
            accept_thread.daemon = True
            accept_thread.start()
            
            return True
        except Exception as e:
            logger.error(f"服务器启动失败: {str(e)}")
            return False
    
    def _accept_connections(self):
        """接受客户端连接的循环"""
        while self.running:
            try:
                client_socket, address = self.server_socket.accept()
                logger.info(f"已连接客户端: {address}")
                
                # 为每个客户端启动一个处理线程
                client_thread = threading.Thread(
                    target=self._handle_client,
                    args=(client_socket, address)
                )
                client_thread.daemon = True
                client_thread.start()
                
                self.clients.append((client_socket, address))
            except Exception as e:
                if self.running:
                    logger.error(f"接受连接时出错: {str(e)}")
                # 如果服务器已停止，就退出循环
                else:
                    break
    
    def _handle_client(self, client_socket: socket.socket, address: tuple):
        """处理单个客户端的通信"""
        try:
            while self.running:
                # 接收数据长度
                length_data = client_socket.recv(4)
                if not length_data:
                    break
                
                message_length = int.from_bytes(length_data, byteorder='big')
                
                # 接收完整消息
                chunks = []
                bytes_received = 0
                while bytes_received < message_length:
                    chunk = client_socket.recv(min(message_length - bytes_received, 4096))
                    if not chunk:
                        raise ConnectionError("连接中断")
                    chunks.append(chunk)
                    bytes_received += len(chunk)
                
                # 组合完整消息
                data = b''.join(chunks)
                
                # 验证和处理消息
                try:
                    message = json.loads(data.decode('utf-8'))
                    
                    # 验证消息签名（如果启用）
                    if self.secret:
                        if not self._verify_message(message):
                            logger.warning(f"来自 {address} 的消息验证失败")
                            self._send_error(client_socket, "验证失败")
                            continue
                    
                    # 处理请求
                    request_data = message.get("data", {})
                    response_data = self.request_handler(request_data)
                    
                    # 发送响应
                    self._send_response(client_socket, response_data)
                    
                except json.JSONDecodeError:
                    logger.error("无效的JSON数据")
                    self._send_error(client_socket, "无效的JSON格式")
                except Exception as e:
                    logger.error(f"处理请求时出错: {str(e)}")
                    self._send_error(client_socket, f"服务器错误: {str(e)}")
        
        except ConnectionError as e:
            logger.info(f"客户端 {address} 断开连接: {str(e)}")
        except Exception as e:
            logger.error(f"处理客户端 {address} 时出错: {str(e)}")
        finally:
            client_socket.close()
            if (client_socket, address) in self.clients:
                self.clients.remove((client_socket, address))
            logger.info(f"已关闭与 {address} 的连接")
    
    def _verify_message(self, message: Dict[str, Any]) -> bool:
        """验证消息的签名"""
        if not self.secret:
            return True
        
        signature = message.get("signature")
        if not signature:
            return False
        
        # 移除签名字段后计算消息的HMAC
        message_copy = message.copy()
        message_copy.pop("signature", None)
        message_str = json.dumps(message_copy, sort_keys=True)
        
        # 计算HMAC
        expected_signature = hmac.new(
            self.secret.encode('utf-8'),
            message_str.encode('utf-8'),
            hashlib.sha256
        ).hexdigest()
        
        return hmac.compare_digest(signature, expected_signature)
    
    def _send_response(self, client_socket: socket.socket, data: Dict[str, Any]):
        """向客户端发送响应"""
        response = {
            "status": "success",
            "data": data
        }
        
        # 添加签名（如果启用）
        if self.secret:
            response_str = json.dumps(response, sort_keys=True)
            signature = hmac.new(
                self.secret.encode('utf-8'),
                response_str.encode('utf-8'),
                hashlib.sha256
            ).hexdigest()
            response["signature"] = signature
        
        # 编码响应
        response_bytes = json.dumps(response).encode('utf-8')
        
        # 发送长度和数据
        length_bytes = len(response_bytes).to_bytes(4, byteorder='big')
        client_socket.sendall(length_bytes + response_bytes)
    
    def _send_error(self, client_socket: socket.socket, error_message: str):
        """向客户端发送错误响应"""
        response = {
            "status": "error",
            "message": error_message
        }
        
        # 添加签名（如果启用）
        if self.secret:
            response_str = json.dumps(response, sort_keys=True)
            signature = hmac.new(
                self.secret.encode('utf-8'),
                response_str.encode('utf-8'),
                hashlib.sha256
            ).hexdigest()
            response["signature"] = signature
        
        # 编码响应
        response_bytes = json.dumps(response).encode('utf-8')
        
        # 发送长度和数据
        length_bytes = len(response_bytes).to_bytes(4, byteorder='big')
        client_socket.sendall(length_bytes + response_bytes)
    
    def stop(self):
        """停止MCP服务器"""
        self.running = False
        
        # 关闭所有客户端连接
        for client_socket, _ in self.clients:
            try:
                client_socket.close()
            except:
                pass
        self.clients = []
        
        # 关闭服务器套接字
        if self.server_socket:
            try:
                self.server_socket.close()
            except:
                pass
        
        logger.info("MCP服务器已停止") 